import pandas as pd
from utils.data_split import DataSplitUtil
from utils.draw_pic import *

def main():
    clinic_data = pd.read_csv(opj(base_path,'data/original_data/Clinic/patient-clinical-data-process.csv'))
    dp = DataSplitUtil(split_random_state=split_random_state_list[0])
    train_raw, test_raw = dp.get_train_test_df(clinic_data)

    # 4. 添加 Dataset 列，用于区分
    train_raw['Dataset'] = 'Train'
    test_raw['Dataset'] = 'Test'

    violin_save_path = opj(clinic_analysis_result_path,'violin')
    box_save_path = opj(clinic_analysis_result_path,'box')
    md(violin_save_path)
    md(box_save_path)

    print(f"数据分布小提琴图已保存在:{opj(clinic_analysis_result_path,'violin')}目录下")
    print(f"数据分布箱线图已保存在:{opj(clinic_analysis_result_path,'box')}目录下")
    df_all = pd.concat([train_raw, test_raw], ignore_index=True)
    draw_violin_by_set(
        df=df_all,
        y_col='Age(years)',
        y_label='Age (years)',
        title='Age Distribution in Train/Test by ALN Status',
        save_path=opj(violin_save_path, 'age.pdf')
    )

    draw_violin_by_set(
        df=df_all,
        y_col='Tumour Size(cm)',
        y_label='Tumor Size (cm)',
        title='Tumor Size Distribution in Train/Test by ALN Status',
        save_path=opj(violin_save_path, 'tumor_size.pdf')
    )

    draw_box_with_centered_points(
        df=df_all,
        y_col='Age(years)',
        y_label='Age (years)',
        title='Age Distribution in Train/Test by ALN Status',
        save_path=opj(box_save_path, 'age.pdf')
    )

    draw_box_with_centered_points(
        df=df_all,
        y_col='Tumour Size(cm)',
        y_label='Tumor Size (cm)',
        title='Tumor Size Distribution in Train/Test by ALN Status',
        save_path=opj(box_save_path, 'tumor_size.pdf')
    )